/*
* basic infer engine for classification tasks
* the pipeline is identical for most models.
*/
#pragma once

#include "image_infer_base.h"

namespace vortex
{
    constexpr float IMAGENET_DEFAULT_MEAN[3] = {0.485, 0.456, 0.406};
    constexpr float IMAGENET_DEFAULT_STD[3] = {0.229, 0.224, 0.225};

    /*
    * image classification with single input/output network
    */
    class classificationEngine : public ImageInferBase
    {
    protected:
        BlobInfo m_InputInfo;
        BlobInfo m_OutputInfo;
    public:
        classificationEngine(
            const std::string& engine_path,
            const BlobInfo& input_info,
            const BlobInfo& output_info
        );

        virtual ~classificationEngine() {}
        // do inference for single image
        virtual bool Infer(const std::string& image_path, std::vector<float>& output);
        
        // do inference for multiple images
        virtual bool InferN(const std::vector<std::string>& image_paths, std::vector<std::vector<float>>& output);
        
        // default preprocess function
        virtual bool Preprocess(const std::string& image_path, int offset = 0);
    };
}
